import os
from typing import List, Tuple
from PyPDF2 import PdfReader
from whoosh import index
from whoosh.fields import Schema, TEXT, ID, NUMERIC
from whoosh.qparser import QueryParser


def find_pdf_files(directory: str) -> List[str]:
    """递归查找目录下所有PDF文件"""
    pdf_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(".pdf"):
                pdf_files.append(os.path.join(root, file))
    return pdf_files


def extract_text_from_pdf(pdf_path: str) -> List[Tuple[int, str]]:
    """提取PDF每一页的文本，返回(页码, 文本)列表"""
    texts = []
    try:
        reader = PdfReader(pdf_path)
        for i, page in enumerate(reader.pages):
            text = page.extract_text() or ""
            texts.append((i + 1, text))
    except Exception as e:
        print(f"提取 {pdf_path} 时出错: {e}")
    return texts


def create_index(index_dir: str, pdf_dir: str):
    """建立Whoosh索引"""
    if not os.path.exists(index_dir):
        os.mkdir(index_dir)
    schema = Schema(
        path=ID(stored=True), page=NUMERIC(stored=True), content=TEXT(stored=True)
    )
    if not index.exists_in(index_dir):
        ix = index.create_in(index_dir, schema)
    else:
        ix = index.open_dir(index_dir)
    writer = ix.writer()
    pdf_files = find_pdf_files(pdf_dir)
    print(f"共找到 {len(pdf_files)} 个PDF文件。")
    for pdf in pdf_files:
        print(f"正在处理: {pdf}")
        pages = extract_text_from_pdf(pdf)
        for page_num, text in pages:
            writer.add_document(path=pdf, page=page_num, content=text)
    writer.commit()
    print("索引建立完成。")


def search_index(index_dir: str, query_str: str, limit: int = 10):
    """搜索Whoosh索引"""
    if not index.exists_in(index_dir):
        print("索引目录不存在，请先建立索引。")
        return
    ix = index.open_dir(index_dir)
    with ix.searcher() as searcher:
        parser = QueryParser("content", schema=ix.schema)
        query = parser.parse(query_str)
        results = searcher.search(query, limit=limit)
        for hit in results:
            print(f"文件: {hit['path']}  页码: {hit['page']}")
            snippet = hit.highlights("content")
            print(f"内容片段: {snippet}\n")


if __name__ == "__main__":
    print("PDF全文索引工具")
    print("1. 建立索引")
    print("2. 搜索")
    choice = input("请选择操作(1/2): ").strip()
    if choice == "1":
        pdf_dir = input("请输入要索引的PDF目录路径：").strip()
        index_dir = input("请输入索引存储目录(如: pdf_index): ").strip()
        create_index(index_dir, pdf_dir)
    elif choice == "2":
        index_dir = input("请输入索引存储目录(如: pdf_index): ").strip()
        while True:
            query_str = input("请输入要搜索的关键词(回车退出): ").strip()
            if not query_str:
                break
            search_index(index_dir, query_str)
    else:
        print("无效选择。")
